package fsm

import (
	"time"

	"gitee.com/hasika/gotool/generic"
)

type StateDuration interface {
	Nanoseconds() int64
}

type IState[State generic.Primitive, Arg any] interface {
	DefaultDuration(Arg) time.Duration
	GetNextStateOnTimeOut(Arg) State
	OnEnter(Arg, bool) StateDuration
	OnExit(Arg, bool)
	Update(time.Duration, Arg)
}

type SwitchStateFunc[State generic.Primitive, Arg any] func(arg Arg, oldState State, newState State)
type ExitStateFunc[State generic.Primitive, Arg any] func(arg Arg, oldState State)

type registry[State generic.Primitive, Arg any] struct {
	defaultState     State
	defaultStateInfo IState[State, Arg]

	stateMap  map[State]IState[State, Arg]
	preEnter  SwitchStateFunc[State, Arg]
	postEnter SwitchStateFunc[State, Arg]
	preExit   ExitStateFunc[State, Arg]
	postExit  SwitchStateFunc[State, Arg]
}

type FSM[State generic.Primitive, Arg any] struct {
	runtimeData[State, Arg]
	registry[State, Arg]
}

func NewFSM[State generic.Primitive, Arg any](defaultState State) *FSM[State, Arg] {
	return &FSM[State, Arg]{
		runtimeData: runtimeData[State, Arg]{
			lastState:    defaultState,
			currentState: defaultState,
		},
		registry: registry[State, Arg]{
			defaultState: defaultState,

			stateMap:  make(map[State]IState[State, Arg]),
			preEnter:  func(arg Arg, oldState State, newState State) {},
			postEnter: func(arg Arg, oldState State, newState State) {},
			preExit:   func(arg Arg, oldState State) {},
			postExit:  func(arg Arg, oldState State, newState State) {},
		},
	}
}

func (f *FSM[State, Arg]) StateLeftDuration() time.Duration {
	return f.stateLeftDuration
}

func (f *FSM[State, Arg]) StateTotalDuration() time.Duration {
	return f.stateTotalDuration
}

func (f *FSM[State, Arg]) Start(arg Arg) {
	f.switchState(f.defaultState, arg, false)
}

func (f *FSM[State, Arg]) Switch(state State, arg Arg) {
	f.switchState(state, arg, false)
}

func (f *FSM[State, Arg]) switchState(nextState State, arg Arg, timeout bool) {
	currentState := f.currentState
	currentStateInfo := f.currentStateInfo
	//老的退出
	f.preExit(arg, currentState)
	if currentStateInfo != nil {
		currentStateInfo.OnExit(arg, timeout)
	}
	nextStateInfo := f.stateMap[nextState]
	f.postExit(arg, currentState, nextState)
	//
	f.lastState = currentState
	//新的进入
	f.preEnter(arg, currentState, nextState)
	f.currentState = nextState
	f.currentStateInfo = nextStateInfo
	defaultDuration := nextStateInfo.DefaultDuration(arg)
	useDuration := defaultDuration
	enterDuration := nextStateInfo.OnEnter(arg, timeout)
	if enterDuration != nil {
		useDuration = time.Duration(enterDuration.Nanoseconds()) * time.Nanosecond
	}
	f.stateTotalDuration = useDuration
	f.stateLeftDuration = useDuration
	f.postEnter(arg, currentState, nextState)
}

func (f *FSM[State, Arg]) Register(state State, iState IState[State, Arg]) {
	f.stateMap[state] = iState
	if state == f.defaultState {
		f.defaultStateInfo = iState
	}
}

func (f *FSM[State, Arg]) Update(arg Arg, dt time.Duration) {
	if f.currentStateInfo == nil {
		f.Start(arg)
	}
	if f.stateLeftDuration != -1 {
		f.stateLeftDuration -= dt
		if f.stateLeftDuration <= 0 {
			nextState := f.currentStateInfo.GetNextStateOnTimeOut(arg)
			f.switchState(nextState, arg, true)
		}
	}
	if f.currentStateInfo != nil {
		f.currentStateInfo.Update(dt, arg)
	}
}

func (f *FSM[State, Arg]) SetPreEnter(preEnter SwitchStateFunc[State, Arg]) {
	f.preEnter = preEnter
}

func (f *FSM[State, Arg]) SetPostEnter(postEnter SwitchStateFunc[State, Arg]) {
	f.postEnter = postEnter
}

func (f *FSM[State, Arg]) SetPreExit(preExit ExitStateFunc[State, Arg]) {
	f.preExit = preExit
}

func (f *FSM[State, Arg]) SetPostExit(postExit SwitchStateFunc[State, Arg]) {
	f.postExit = postExit
}

func (f *FSM[State, Arg]) CurrentState() State {
	return f.currentState
}

func (f *FSM[State, Arg]) SetStateLeftDuration(duration time.Duration) {
	f.stateLeftDuration = duration
}
